/**
 * 
 */
package edu.umd.clip.lm.storage.remote;

import java.io.*;
import java.lang.reflect.*;
import java.nio.*;
import java.nio.channels.*;
import java.util.concurrent.atomic.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
abstract public class Message {
	public static final byte CURRENT_VERSION = 1;
	
	@SuppressWarnings("unchecked")
	protected static class MessageHeader {
		public static final byte SINGLE_TREE_REQUEST = 1;
		public static final byte SINGLE_TREE_RESPONSE = 2;
		public static final byte FAILURE_RESPONSE = 3;
		public static final byte PING_MESSAGE = 4;
		
		private static final Class<? extends Message> msgTypes[];
		static {
			msgTypes = new Class[5];
			msgTypes[SINGLE_TREE_REQUEST] = SingleTreeRequestMessage.class;
			msgTypes[SINGLE_TREE_RESPONSE] = SingleTreeResponseMessage.class;
			msgTypes[FAILURE_RESPONSE] = FailureResponseMessage.class;
			msgTypes[PING_MESSAGE] = PingMessage.class;
		}
		
		public byte version;
		public byte type;
		public int id;
		public int modelId;
		public int replyTo = -1;
		public int length;
		
		public static final int HEADER_SIZE = 1 + 1 + 4 + 4 + 4 + 4;
		
		private MessageHeader(byte type, int id, int modelId) {
			this(type, id, modelId, -1);
		}
		
		private MessageHeader(byte type, int id, int modelId, int length) {
			this.version = CURRENT_VERSION;
			this.type = type;
			this.id = id;
			this.modelId = modelId;
			this.length = length;
		}
		
		public MessageHeader(ByteBuffer buffer) {
			version = buffer.get();
			type = buffer.get();
			id = buffer.getInt();
			modelId = buffer.getInt();
			replyTo = buffer.getInt();
			length = buffer.getInt();
		}
		
		public void write(ByteBuffer buffer) {
			buffer.put(version);
			buffer.put(type);
			buffer.putInt(id);
			buffer.putInt(modelId);
			buffer.putInt(replyTo);
			buffer.putInt(length);
		}
	}
	
	private MessageHeader header;
	private static AtomicInteger idSource = new AtomicInteger(0);
	
	protected Message(byte type, int modelId, int length) {
		int id = idSource.incrementAndGet();
		header = new MessageHeader(type, id, modelId, length);
	}
	
	protected Message(byte type, int modelId) {
		int id = idSource.incrementAndGet();
		header = new MessageHeader(type, id, modelId);
	}
	
	protected Message(MessageHeader header) {
		this.header = header;
	}
	
	public static Message getMessage(ReadableByteChannel channel) throws IOException {
		ByteBuffer buffer = ByteBuffer.allocate(MessageHeader.HEADER_SIZE);
		readFully(channel, buffer);
		if (buffer.position() == 0) {
			// asynchronous channel, no data
			return null;
		}
		buffer.rewind();
		MessageHeader header = new MessageHeader(buffer);
		
		buffer = ByteBuffer.allocate(header.length);
		readFully(channel, buffer);
		buffer.rewind();
		
		if (header.type >= MessageHeader.msgTypes.length || MessageHeader.msgTypes[header.type] == null) {
			throw new Error("unknown/unimplemented message type: " + Byte.toString(header.type));
		}
		Class<? extends Message> cl = MessageHeader.msgTypes[header.type];
		try {
			Constructor<? extends Message> constructor = cl.getConstructor(MessageHeader.class, ByteBuffer.class);
			return constructor.newInstance(header, buffer);
		} catch (Exception e) {
			e.printStackTrace();
			throw new Error("failed to construct a message");
		}
	}
	
	public ByteBuffer write() {
		ByteBuffer bodyBuffer = writeBody();
		int bodySize = bodyBuffer.limit();
		ByteBuffer msgBuffer = ByteBuffer.allocate(MessageHeader.HEADER_SIZE + bodySize);
		header.length = bodySize; 
		header.write(msgBuffer);
		bodyBuffer.rewind();
		msgBuffer.put(bodyBuffer);
		msgBuffer.rewind();
		return msgBuffer;
	}
	
	public int write(ByteBuffer buffer) {
		ByteBuffer bodyBuffer = writeBody();
		int bodySize = bodyBuffer.limit();
		header.length = bodySize; 
		header.write(buffer);
		bodyBuffer.rewind();
		buffer.put(bodyBuffer);
		
		return bodySize + MessageHeader.HEADER_SIZE;
	}
	
	public int send(WritableByteChannel channel) throws IOException {
		ByteBuffer msgBuffer = write();
		channel.write(msgBuffer);
		
		if (msgBuffer.remaining() > 0) {
			throw new IOException("failed to write a message!");
		}
		return msgBuffer.limit();
	}
	
	public int getId() {
		return header.id;
	}
	
	public int getReplyTo() {
		return header.replyTo;
	}
	
	protected void setReplyTo(int replyTo) {
		header.replyTo = replyTo;
	}

	public int getModelId() {
		return header.modelId;
	}

	public byte getType() {
		return header.type;
	}
	
	public int getSize() {
		return header.length + MessageHeader.HEADER_SIZE;
	}
	
	protected abstract void readBody(ByteBuffer buffer);
	protected abstract ByteBuffer writeBody();
	
	// a helper function
	private static void readFully(ReadableByteChannel channel, ByteBuffer buffer) throws IOException {
		int read = 0;
		do {
			int size = channel.read(buffer);
			if (size == -1) throw new IOException("unexpected end of input");
			read += size;
			if (read == 0) {
				//System.out.print("");
				return;
			}
		} while(read < buffer.capacity());		
	}
}
